# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
from typing import Union, List, Any
from fastapi import Query
from pydantic import BaseModel, Field


class KnowArchiveWholeIn(BaseModel):
    """ 全部参数 """
    kb: str = Query(..., min_length=1, max_length=32, description="知识库编号")


class KnowArchiveSearchIn(BaseModel):
    """ 文档搜索参数 """
    page_no: int = Query(default=1, gt=0, description="当前页码")
    page_size: int = Query(default=15, gt=0, le=200, description="每页条数")
    keyword: Union[str, None] = Query(default=None, description="关键词")
    kb: str = Query(..., min_length=1, max_length=32, description="知识库编号")


class KnowArchiveDetailIn(BaseModel):
    """ 文档详情参数 """
    fid: str = Query(..., min_length=1, description="文档编号")


class KnowArchiveInspectIn(BaseModel):
    """ 文档详情参数 """
    ids: List[str] = Field(default=[], description="文档编号")


class KnowArchiveRenameIn(BaseModel):
    """ 文档命名参数 """
    fid: str = Field(..., min_length=1, max_length=32, description="文档编号")
    name: str = Field(..., min_length=1, max_length=100, description="名称")

    class Config:
        json_schema_extra = {
            "example": {
                "fid": "70efdf2ec9b086079795c442636b55fb",
                "name": "HCxiS50f0d4nd"
            }
        }


class KnowArchiveMergeIn(BaseModel):
    """ 文档合并参数 """
    source_fid: str = Field(..., min_length=1, max_length=32, description="来源ID")
    target_fid: str = Field(..., min_length=1, max_length=32, description="目标ID")

    class Config:
        json_schema_extra = {
            "example": {
                "source_fid": "9bf8d27aed419a085d56ae125c85968f",
                "target_fid": "75f8627a89415a085052a5125c839686"
            }
        }


class KnowArchiveFidIn(BaseModel):
    """ 文档删除参数 """
    fid: str = Field(..., min_length=1, max_length=32, description="文档编号")

    class Config:
        json_schema_extra = {
            "example": {
                "fid": "9bf8d27aed419a085d56ae125c85968f"
            }
        }


class KnowArchiveImportIn(BaseModel):
    """ 文档导入参数 """
    kb: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    scene: str = Field(..., min_length=1, pattern=r"^(file|text|link|csv)$", description="解析场景")
    train_mode: str = Field(..., min_length=1, pattern=r"^(chunk|qa)$", description="训练模式")
    chunk_size: int = Field(default=512, ge=10, le=32000, description="分段长度")
    chunk_overlap: int = Field(default=0, ge=0, le=1000, description="重叠长度")
    separators: str = Field(default="", max_length=200, description="分段规则")
    files: list = Field(..., description="文件地址")

    class Config:
        json_schema_extra = {
            "example": {
                "kb": "RuMiVZPMoc5ih3aj",
                "scene": "file",
                "train_mode": "chunk",
                "chunk_size": 512,
                "chunk_overlap": 0,
                "separators": "",
                "files": [
                    {
                        "name": "20230625O.docx",
                        "url": "https://xx.com/ab30d941298.docx"
                    }
                ]
            }
        }


class KnowArchiveParseIn(BaseModel):
    """ 文档解析参数 """
    kb: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    scene: str = Field(..., min_length=1, pattern=r"^(file|text|link|csv)$", description="解析场景")
    chunk_size: int = Field(default=512, ge=10, le=32000, description="分段长度")
    chunk_overlap: int = Field(default=0, ge=0, le=1000, description="重叠长度")
    separators: str = Field(default="", max_length=200, description="分段规则")
    url: str = Field(..., min_length=1, description="文件地址")

    class Config:
        json_schema_extra = {
            "example": {
                "kb": "RuMiVZPMoc5ih3aj",
                "scene": "file",
                "chunk_size": 512,
                "chunk_overlap": 0,
                "separators": "",
                "url": "https://xx.com/ab30d941298.docx"
            }
        }


"""--------------- Separator ---------------"""


class KnowArchiveWholeVo(BaseModel):
    """ 全部文档VO """
    code: str = Field(description="文件编号")
    name: str = Field(description="文件名称")


class KnowArchiveListVo(BaseModel):
    """ 文档列表VO """
    code: str = Field(description="文件编号")
    name: str = Field(description="文件名称")
    scene: str = Field(description="来源场景")
    qa_status: int = Field(description="拆分状态: [0=待处理, 1=排队中, 2=拆分中, 3=成功, 4=失败]")
    train_mode: str = Field(description="训练模式: [chunk=直接分段, qa=问答拆分]")
    train_total: int = Field(description="数据总量")
    train_waits: int = Field(description="等待训练")
    fail_reason: str = Field(description="失败原因")
    is_disable: int = Field(description="是否禁用: [0=否, 1=是]")
    create_time: str = Field(description="创建时间")
    update_time: str = Field(description="更新时间")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "70efdf2ec9b086079795c442636b55fb",
                "name": "test.txt",
                "scene": "file",
                "qa_status": 0,
                "train_mode": "chunk",
                "train_total": 1,
                "train_waits": 0,
                "is_disable": 0,
                "fail_reason": "失败原因",
                "create_time": "2024-09-28 17:16:48",
                "update_time": "2024-09-28 17:16:48"
            }
        }


class KnowArchiveDetailVo(BaseModel):
    """ 文档详情VO """
    code: str = Field(description="文档编号")
    name: str = Field(description="文档名称")
    size: int = Field(description="文档大小(KB)")
    train_mode: str = Field(description="训练模式: [chunk=直接分段, qa=问答拆分]")
    is_disable: int = Field(description="是否禁用: [0=否, 1=是]")
    create_time: str = Field(description="创建时间")
    update_time: str = Field(description="更新时间")

    class Config:
        json_schema_extra = {
            "example": {
                "code": "c9f0f895fb98ab9159f51fd0297e236d",
                "name": "test.docx",
                "size": 100,
                "train_mode": "chunk",
                "is_disable": 0,
                "create_time": "2024-08-27 14:40:21",
                "update_time": "2024-08-27 14:40:21"
            }
        }


class KnowArchiveParseVo(BaseModel):
    """ 文档解析VO """
    split_count: int = Field(description="段落数量")
    lists: List[Any] = Field(description="拆分列表")

    class Config:
        json_schema_extra = {
            "example": {
                "split_count": 24,
                "lists": [
                    {
                        "question": "开户行：中国工商银行广州经济技术开发区支行"
                    }
                ]
            }
        }
